…nd KV cache utils

Step 1 (SpeculativeConfig):
- Fallback to dflash_config.target_layer_ids in compute_hash() when
  eagle_aux_hidden_state_layer_ids is not set, with +1 shift to match
  vLLM hidden-state extraction semantics
- New requires_eagle_cache_drop() property that returns False for DFlash
  (DFlash writes all context KV before drafting, so no cache drop needed)

Step 2 (speculators/algos.py):
- Forward SWA-related config keys (layer_types, use_sliding_window,
  sliding_window, max_window_layers) through update_dflash()
- Shift eagle_aux_hidden_state_layer_ids by +1 to align with vLLM's
  layer indexing (draft model uses 0-based, runner uses 1-based)
- Remove TODO comment now that the shift is applied

Step 3 (scheduler.py):
- Add self.requires_eagle_cache_drop attribute derived from
  speculative_config.requires_eagle_cache_drop()
- Replace self.use_eagle with self.requires_eagle_cache_drop in:
  KVCacheManager constructor, mamba block-aligned split guard, and
  encoder input scheduling (2 call sites). DFlash speculative decoding
  no longer triggers unnecessary cache pruning.

Step 4 (kv_cache_utils.py):
- Replace index-based loop with named variable over kv_cache_groups for
  readability and safety

Co-authored-by: vLLM upstream PR vllm-project#40898 contributors

…support in qwen3_dflash

Changes to vllm/model_executor/models/qwen3_dflash.py:

- New imports: Mapping (for per-layer slot mapping), FullAttentionSpec,
  KVCacheSpec, SlidingWindowSpec (for KV cache spec widening)

- DFlashAttention class (subclass of Attention): overrides get_kv_cache_spec()
  to widen SlidingWindowSpec to FullAttentionSpec. DFlash writes every context
  KV before drafting and cannot evict old context blocks from sliding-window
  layers, so the KV cache must be allocated as full attention.

- _get_dflash_layer_types() helper: resolves per-layer attention type from
  config.layer_types, defaults to all full_attention. Validates layer type
  names and that sliding_window is configured when sliding_attention is used.

- DFlashQwen3DecoderLayer: added layer_type parameter, tracked on self.
  Kept sliding_window uniform across all layers (getattr from config) to
  maintain single KV cache group for the drafter model.

- DFlashQwen3Model: uses layer_types to configure decoder layers, exposes
  sliding_attention_layer_names set for the proposer's metadata building.

- precompute_and_store_context_kv: context_slot_mapping now accepts
  Mapping[str, torch.Tensor] for per-layer slot assignments. Cache insert
  loop extracts the correct slot mapping per layer.

- DFlashQwen3ForCausalLM: updated signature to match, exposes
  sliding_attention_layer_names property.

Note: The original code already had per_layer_sliding_window=sliding_window
in Attention(). The conditional layer_type-specific sliding_window from the
upstream PR was NOT applied because it would split DFlash layers into
different KV cache groups, breaking the single-group assertion in
llm_base_proposer.py. Instead, the layer_type tracking is purely for the
proposer's per-layer causal/non-causal metadata building.

Co-authored-by: vLLM upstream PR vllm-project#40898 contributors

…n base proposer

Changes to vllm/v1/spec_decode/llm_base_proposer.py:

- New imports: KVCacheSpec (added to existing KVCacheConfig import)

- New attributes on SpecDecodeBaseProposer:
  _draft_layer_to_kv_cache_gid: maps each draft layer name to its KV cache
  group ID
  _draft_kv_cache_group_ids: sorted list of unique group IDs used by
  draft layers

- New hook method allow_multiple_draft_kv_cache_groups() -> bool:
  Returns False by default (single-group constraint for EAGLE/draft models).
  DFlashProposer overrides this to return True.

- Rewrite initialize_attn_backend(): Replaces the monolithic single-group
  lookup with per-layer-to-group mapping. The attention_groups dict key now
  includes (gid, backend_key, layer_kv_cache_spec) instead of just
  (backend_key), allowing draft layers from different KV cache groups to
  be properly grouped. The single-group validation is now gated behind
  allow_multiple_draft_kv_cache_groups().

Backend-agnostic: Uses abstract get_attn_backend(), AttentionGroup, and
KVCacheSpec interfaces. No backend-specific code.

Co-authored-by: vLLM upstream PR vllm-project#40898 contributors

…rs, slot mappings, SWA metadata

Changes to vllm/v1/spec_decode/dflash.py:

Imports:
- Moved replace from dataclasses to vllm.config (VllmConfig, replace)
- Added KVCacheConfig import

New __init__ attributes:
- _slot_mapping_buffers_by_gid: per-KV-group (context, query) slot mapping pairs
- _draft_block_size_by_gid: per-KV-group block sizes for triton kernel
- _draft_block_tables: per-KV-group block tables (set by gpu_model_runner)

New override methods:
- allow_multiple_draft_kv_cache_groups() -> True: enables multi-KV-group support
- initialize_attn_backend(): calls super() then populates per-KV-group block
  sizes and ensures slot mapping buffers are allocated
- clear_draft_block_tables(): resets block tables before each step
- set_draft_block_table(kv_cache_gid, block_table): receives block tables from
  gpu_model_runner per KV cache group

New helper methods:
- _ensure_slot_mapping_buffers(): lazy-allocates per-KV-group buffers,
  reuses existing buffers for the primary group
- _draft_kv_gids(): resolves draft KV group IDs from inherited attributes
- _get_dflash_block_table(kv_cache_gid, cad): looks up per-KV-group block
  table, falls back to cad's block table
- _get_dflash_context_slot_mapping(num_context): returns context slot mappings
  as dict[str, torch.Tensor] when layers span multiple KV groups
- _get_slot_mapping() override: returns per-layer slot mappings when
  layers span multiple KV groups

Rewrote set_inputs_first_pass():
- Loops over draft KV group IDs, dispatching triton kernel per group
- Uses per-KV-group block tables, slot mapping buffers, and block sizes
- Builds new CommonAttentionMetadata with primary group's block table

Updated build_model_inputs_first_pass():
- Calls _get_dflash_context_slot_mapping() for per-layer context slot mapping

Rewrote build_per_group_and_layer_attn_metadata():
- Builds per-group metadata with per-KV-group block tables and slot mappings
- Applies causal metadata overlay for SWA layers (sliding_attention_layer_names)
- Asserts causal=True for SWA layers, causal=False for full attention layers

FlashInfer adaptation note:
- build_for_drafting(causal=True/False) works correctly with FlashInfer because
  FlashInferMetadataBuilder.build() reads common_attn_metadata.causal and
  FlashInferMetadata has causal: bool field. The per-layer causal override for
  SWA layers is handled by the causal/non-causal assertion checks.

Co-authored-by: vLLM upstream PR vllm-project#40898 contributors

…dge + DFlash block table injection

Changes to vllm/v1/worker/gpu_model_runner.py:

Change A - DFlash block table injection:
- In the attention metadata build loop, detect DFlashProposer and call
  clear_draft_block_tables() before the loop
- Inside the kv_cache_gid loop, call set_draft_block_table() to inject
  per-KV-group block tables into the DFlash drafter

Change C - New helper methods for KV cache stride/mapping:
- _get_kv_cache_stride_order(): Extracts stride order from AttentionBackend
  with fallback to identity ordering
- _get_standard_kv_cache_orders(): Maps backend stride order to named
  public/physical dimension orderings (kv/block/token/head/dim), detecting
  (2, num_blocks, ...) vs (num_blocks, 2, ...) layouts
- _view_kv_cache_with_physical_order(): Creates torch.as_strided view with
  physical stride ordering for cross-backend stride compatibility
- _get_attention_kv_cache_shape(): Computes KV cache shape from spec,
  handling storage_block_size != block_size for MLA with compression
- _get_raw_tensor_physical_orders(): Scans all attention groups to collect
  physical stride orders per raw tensor, enabling multi-group physical order
  detection

Change D - Rewrite _reshape_kv_cache_tensors attention block:
- Uses new helper methods instead of inline shape/stride computation
- Adds physical-order bridging: when a raw tensor is shared across multiple
  attention groups with different stride orders, creates a strided view
  matching the shared physical order instead of the default backend order
- Falls back to existing contiguous/view/permute path when physical order
  bridging is not applicable (padded pages, mismatched layouts, etc.)

FlashInfer adaptation:
- Helper methods call get_kv_cache_stride_order() and get_kv_cache_shape()
  dynamically on the backend type — works identically for FlashInfer
- Physical-order bridging uses torch.as_strided for stride remapping,
  independent of kernel-level layout
- No behavioral changes for single-backend (no shared raw tensors)

Co-authored-by: vLLM upstream PR vllm-project#40898 contributors

New test file tests/v1/spec_decode/test_dflash_swa.py (5 tests):
- test_dflash_speculators_preserves_swa_config: SWA fields (layer_types,
  use_sliding_window, sliding_window, max_window_layers) flow through
  SpeculatorsConfig.extract_transformers_pre_trained_config
- test_dflash_compile_hash_uses_checkpoint_layer_id_semantics: Hash
  consistency between dflash_config.target_layer_ids and
  eagle_aux_hidden_state_layer_ids (shifted by +1)
- test_dflash_swa_layers_use_full_kv_cache_spec: DFlashAttention
  get_kv_cache_spec widens SlidingWindowSpec to FullAttentionSpec
- test_dflash_swa_layers_use_causal_metadata:
  build_per_group_and_layer_attn_metadata sets causal=True for SWA
  layers, causal=False for full attention layers
- test_dflash_metadata_uses_per_kv_group_slot_mapping: Per-KV-group
  block tables and slot mappings are correctly assigned to layers

Added to tests/v1/core/test_kv_sharing.py:
- test_dflash_draft_kv_groups_keep_hybrid_tensor_sharing: DFlash with
  multiple KV cache groups keeps tensors shared across groups

Added to tests/v1/worker/test_gpu_model_runner.py:
- test_kv_major_cache_can_share_block_major_raw_tensor: Verifies
  _view_kv_cache_with_physical_order creates correct strided views
  for two different physical orderings

All tests use mock objects (_FakeBuilder, _FakeAttentionGroup) and
SimpleNamespace stubs — fully backend-agnostic, no real GPU or
attention backend required.

All 8 tests pass.

Co-authored-by: vLLM upstream PR vllm-project#40898 contributors